import os
import shutil
import glob

# 获取当前目录
current_dir = os.getcwd()

folders = [
    'train',
    'test'
]

# 创建 test 目录
try:
    os.mkdir('test')
    print("成功创建 test 目录")
except FileExistsError:
    print("test 目录已存在")

# 创建 train 目录
try:
    os.mkdir('train')
    print("成功创建 train 目录")
except FileExistsError:
    print("train 目录已存在")

# 定义要复制的目录模式和目标目录的映射
patterns_and_destinations = [
    ('GenericMOT_JPEG_Sequence/*0/', 'train/'),
    ('GenericMOT_JPEG_Sequence/*1/', 'train/'),
    ('GenericMOT_JPEG_Sequence/*2/', 'train/'),
    ('GenericMOT_JPEG_Sequence/*3/', 'test/')
]

# 遍历每个模式和目标目录的组合
for pattern, destination in patterns_and_destinations:
    # 使用 glob 查找符合模式的目录
    for source_dir in glob.glob(pattern):
        try:
            # 获取源目录的基名
            base_name = os.path.basename(os.path.normpath(source_dir))
            # 构建目标路径
            target_path = os.path.join(destination, base_name)
            # 复制目录
            shutil.copytree(source_dir, target_path)
            print(f"成功将 {source_dir} 复制到 {target_path}")
        except FileNotFoundError:
            print(f"未找到目录 {source_dir}")
        except FileExistsError:
            print(f"目标路径 {target_path} 已存在，跳过复制")
        except Exception as e:
            print(f"复制 {source_dir} 时出现错误: {e}")


for folder in folders:
    # 定义 test 文件夹路径
    src_folder = os.path.join(current_dir, folder)

    # 定义上级目录下的 track_label 文件夹路径
    # parent_dir = os.path.dirname(current_dir)
    track_label_folder = os.path.join(current_dir, 'track_label')

    # 检查 test 文件夹是否存在
    if not os.path.exists(src_folder):
        print(f"{src_folder} 文件夹不存在。")
    else:
        # 遍历 test 文件夹下的一级子目录
        for sub_dir_name in os.listdir(src_folder):
            sub_dir_path = os.path.join(src_folder, sub_dir_name)
            if os.path.isdir(sub_dir_path):
                # 在一级子目录下创建 gt 文件夹，如果已存在则不创建
                gt_folder = os.path.join(sub_dir_path, 'gt')
                if not os.path.exists(gt_folder):
                    try:
                        os.makedirs(gt_folder)
                        print(f"在 {sub_dir_path} 下成功创建 gt 文件夹。")
                    except FileExistsError:
                        pass
                    except Exception as e:
                        print(f"在 {sub_dir_path} 下创建 gt 文件夹时出现错误: {e}")
                else:
                    print(f"{gt_folder} 已存在，无需创建。")

                # 查找与子目录同名的 txt 文件并复制到 gt 文件夹中
                txt_file_name = sub_dir_name + '.txt'
                txt_file_path = os.path.join(track_label_folder, txt_file_name)
                destination_file_path = os.path.join(gt_folder, 'gt.txt')

                if os.path.exists(txt_file_path):
                    if not os.path.exists(destination_file_path):
                        try:
                            shutil.copy2(txt_file_path, destination_file_path)
                            print(f"已将 {txt_file_path} 复制到 {destination_file_path}。")
                        except Exception as e:
                            print(f"复制 {txt_file_path} 到 {destination_file_path} 时出现错误: {e}")
                    else:
                        print(f"{destination_file_path} 已存在，无需复制。")
                else:
                    print(f"{txt_file_path} 不存在，无法复制。")